# 多参数估计数据集Agent
import numpy as np
import torch
from apps.fmcw.conf.app_config import AppConfig as AF
from apps.fmcw.core.ira_observe import IraObserve
from apps.fmcw.core.ira_action import IraAction
from apps.fmcw.core.fmcw_rsp import FmcwRsp as RSP

class MpeDssAgent(object):
    IDX = 0

    def __init__(self):
        self.name = 'apps.fmcw.core.mpe_dss_agent.MpeDssAgent'

    def notify(self, obs:IraObserve) -> IraAction:
        action = IraAction(n=3)
        action.dt_action[0] = 0
        RDC = obs.RDC
        ranges_, velocitys, thetas = AF.ranges_, AF.velocitys, AF.thetas
        X = torch.from_numpy(RDC)
        print(f'??? X: {type(X)}, {X.shape};')
        y = torch.zeros((AF.MAX_OBJS*AF.OBJ_DIMS), dtype=torch.float32)
        for idx in range(len(ranges_)):
            if ranges_[idx] < 1.0 or ranges_[idx] > AF.Rmax-1:
                action.dt_action[0] = 888 # 魔术数字使程序退出
                exit(0)
            y[idx*AF.OBJ_DIMS+0], y[idx*AF.OBJ_DIMS+1], y[idx*AF.OBJ_DIMS+2], y[idx*AF.OBJ_DIMS+3] = 1.0, ranges_[idx], velocitys[idx], thetas[idx]
        print(f'    y: {type(y)}; {y.shape}; \n{y};')
        torch.save(X, f'./work/datasets/mpe/X_{MpeDssAgent.IDX:05}.pt')
        torch.save(y, f'./work/datasets/mpe/y_{MpeDssAgent.IDX:05}.pt')
        MpeDssAgent.IDX += 1
        return action

    @staticmethod
    def verify_dss() -> None:
        ''' 
        验证保存数据集的准确性，将RDC读出来，利用传统信号处理方法进行处理，看结果是否正确
        '''
        idx = 200
        X = torch.load(f'./work/datasets/mpe/X_{idx:05}.pt')
        y = torch.load(f'./work/datasets/mpe/y_{idx:05}.pt')
        print(f'X: {type(X)}; {X.shape};')
        print(f'y: {type(y)}; {y.shape};\n{y};')
        RDC = X.numpy()
        print(f'RDC: {type(RDC)}; {RDC.shape};')
        # 准备环境
        delta = (AF.T * AF.numChirps) / (AF.N_CPI - 1)
        AF.t_sys = AF.T * AF.numChirps * idx + delta
        AF.times = np.linspace(AF.t_sys, AF.t_sys + AF.T * AF.numChirps, AF.N_CPI)
        AF.t_onePulse = np.arange(AF.t_sys, AF.t_sys + AF.dt * AF.numADC, AF.dt)[:256]
        #　雷达信号处理－距离多普勒
        RDMs = RSP.range_doppler(RDC)
        RSP.draw_range_doppler()
        # 雷达信号处理－CA-CFAR
        RDM_mask, cfar_ranges, cfar_dopps, K = RSP.ca_cfar(RDMs)
        RSP.draw_ca_cfar(RDM_mask)
        # 雷达信号处理－MUSIC角度估计
        music_spectrum, a1 = RSP.music_angle(RDMs, cfar_ranges, cfar_dopps, K)
        RSP.draw_music_angle(music_spectrum, K)
        # 雷达信号处理－距离角度
        range_az_music = RSP.range_angle(RDC, a1)
        RSP.draw_range_angle(range_az_music)
        # 雷达信号处理－生成点云
        coor1, coor2 = RSP.point_cloud(music_spectrum, cfar_ranges)
        RSP.draw_point_cloud(coor1, coor2)